import torch.nn as nn

from mmdet.models import NECKS
from models.backbones.mobilenetv2 import InvertedResidual


@NECKS.register_module()
class SSDLiteNeckMobileNetv2(nn.Module):
    def __init__(self, in_channel: int, extra_channels: list):
        super().__init__()
        assert len(extra_channels) == 4
        self.extras = nn.ModuleList([
            InvertedResidual(in_channel,
                             extra_channels[0],
                             stride=2,
                             expand_ratio=0.2),
            InvertedResidual(extra_channels[0],
                             extra_channels[1],
                             stride=2,
                             expand_ratio=0.25),
            InvertedResidual(extra_channels[1],
                             extra_channels[2],
                             stride=2,
                             expand_ratio=0.5),
            InvertedResidual(extra_channels[2],
                             extra_channels[3],
                             stride=2,
                             expand_ratio=0.25),
        ])

    def forward(self, x):
        outs = list()
        for extra in self.extras:
            x = extra(x)
            outs.append(x)
        return outs
